import numpy as np
from typing import Union
import platgo as pg


class Population:
    """
    种群类
    """
    
    def __init__(self, N: int = None, problem: pg.Problem = None, decs: np.ndarray = None, cv: np.ndarray = None,
                 objv: np.ndarray = None) -> None:
        """
        种群初始化。种群需要在算法模板里做初始化，如果没有传入decs，则会调用_init_decs自动根据问题要求初始化种群的决策矩阵
        :param N:        种群大小
        :param problem:  优化问题
        :param decs:     决策矩阵
        :param cv:       约束违反矩阵
        :param objv:     目标值矩阵
        """
        # 矩阵必须是2维
        assert decs is None or decs.ndim == 2, "decs must be 2-D array"
        assert cv is None or cv.ndim == 2, "cv must be 2-D array"
        assert objv is None or objv.ndim == 2, "objv must be 2-D array"
        assert not problem or not decs, "At least one of problem and decs is not None"  # problem和decs至少有一个不为None，否则没有办法初始化决策矩阵
        assert not N or not decs, "At least one of N and decs is not None"  # N和decs至少有一个不为None， 否则不能确定种群大小
        if decs is None:
            self._init_decs(N=N, problem=problem)  # 使用内部方法初始化决策矩阵
        else:
            self.decs = decs.copy()
        self.cv = cv.copy() if cv is not None else None
        self.objv = objv.copy() if objv is not None else None
    
    def _init_decs(self, N: int, problem: pg.Problem) -> None:
        """
        根据问题的要求初始化种群染色体
        :param N: 种群大小，即决策矩阵大小
        :param problem:  根据问题的上下界和要求的编码方式
        """
        lb = problem.borders[0] # 获取问题的取值下界
        ub = problem.borders[1] # 获取问题的取值上界
        lb = np.tile(lb, (N, 1)) # 将lb矩阵向0维坐标方向重复N次
        ub = np.tile(ub, (N, 1)) # 将ub矩阵向0维坐标方向重复N次
        self.decs = np.random.uniform(lb, ub)  # TODO 这里种群初始化的时候要考虑编码方式
    
    @property
    def N(self) -> int:
        return self.decs.shape[0]
    
    def fix_decs(self, borders: np.ndarray, method: int = 0) -> None:
        """
        修复越界
        :param method: 多种修复方法
        :param borders: [lb, ub] 决策变量的上下界
        :return: 修复种群，不需要返回值
        """
        # TODO 对border做检查
        if method == 0:  # 截断修复
            self.decs = np.clip(self.decs, borders[0], borders[1])
        
        # TODO 循环修复 往复修复 随机修复
    
    def copy(self):
        """返回对象副本"""
        new_decs = self.decs.copy()
        new_cv = self.cv.copy() if self.cv is not None else None
        new_objv = self.objv.copy() if self.objv is not None else None
        pop = Population(decs=new_decs, cv=new_cv, objv=new_objv)
        return pop
    
    def __getitem__(self, ind: Union[int, list, np.ndarray, slice]):
        """
        种群切片，根据下标选择部分个体生成新的种群
        ndarray索引分为int下标索引和bool索引，计算N的方式不同
        :param ind: 新种群的索引，接受int, list, ndarray, slice
        """
        if self.decs is None:
            raise RuntimeError('The population has not been initialized')
        if type(ind) == int:
            ind = [ind]
        if type(ind) == np.ndarray:
            # 索引的类型只能是int32或bool
            assert ind.dtype in [np.int32, np.bool]
            # 索引的维度只能是 (n,) 或是 (1,n)
            assert ind.ndim == 1 or ind.ndim == 2
            if ind.ndim == 2:
                assert 1 in ind.shape
                ind = ind.flatten()
        
        new_decs = self.decs[ind]
        new_cv = self.cv[ind] if self.cv is not None else None
        new_objv = self.objv[ind] if self.objv is not None else None
        new_pop = Population(decs=new_decs, cv=new_cv, objv=new_objv)
        return new_pop
    
    def __setitem__(self, item: Union[int, list, np.ndarray, slice], pop) -> None:
        """
        为种群内的部分个体赋值，多对一
        population[[0,1]] = pop
        :param item: 下标
        :param pop: instance of Population
        :return:
        """
        # TODO 两个种群需要进行检查
        if self.decs:
            self.decs[item] = pop.decs
        if self.cv:
            self.cv[item] = pop.cv
        if self.objv:
            self.objv[item] = pop.objv
    
    def __add__(self, pop):
        """
        合并种群,不更改原来的两种群，而是返回新的种群
        不会重新计算目标函数值和约束
        :param pop:
        :return:
        """
        # TODO cv和objv可以是空，但是两种群必须一致
        new_decs = np.vstack([self.decs, pop.decs])
        new_cv = np.vstack([self.cv, pop.cv]) if self.cv is not None else None
        new_objv = np.vstack([self.objv, pop.objv]) if self.objv is not None else None
        new_pop = Population(decs=new_decs, cv=new_cv, objv=new_objv)
        return new_pop
    
    def __len__(self):
        return self.decs.shape[0]
    
    def __str__(self) -> str:
        """打印种群信息"""
        return (
            """
            population information:
            ----------------------
            decs shape = {},
            best objective = {}
            """.format(
                self.decs.shape,
                np.min(self.objv, axis=0, keepdims=True)
            )
        )
